#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@author: iori.king
"""

from flask import Flask, render_template, request, redirect, send_file, flash, jsonify, send_from_directory, make_response
from flask_cors import CORS
from werkzeug.utils import secure_filename
from flask_session import Session
import os
import cv2
import uuid
from PIL import Image
import numpy as np
import base64

sess = Session()
app = Flask(__name__, static_url_path='')

UPLOAD_FOLDER = './uploads'
ALLOWED_EXTENSIONS = {'xml'}

UPLOAD_SUBDIR = os.path.join(os.path.dirname(
    os.path.realpath(__file__)), 'uploads', 'inprocess')
if not os.path.exists(UPLOAD_SUBDIR):
    try:
        os.makedirs(UPLOAD_SUBDIR)
    except OSError as e:
        print(e)


@app.route('/upload', methods=['POST'])
def upload_file():
    # try:
    # if os.path.exists("./folder/new.xml"):
    #    print("removed")
    #    os.remove("./folder/new.xml")
    f = request.files['json']
    f.save(os.path.join(
        app.config['UPLOAD_FOLDER'], secure_filename(f.filename)))
    message = " uploaded successfully"
    return f.filename+message


@app.route('/', methods=['GET'])
def home():
    return send_from_directory('static', 'index.html')


# 抠图，传主文件名
@app.route('/kou', methods=['POST'])
def kou():
    print(request)
    try:
        filename = request.form['filename']
        path = os.path.join(UPLOAD_SUBDIR, filename)
        f_name, ex_name = os.path.splitext(filename)
        if len(request.files) > 0:
            mask_file = request.files['mask']
            mask_ame = mask_file.filename
            mask_path = os.path.join(UPLOAD_SUBDIR, f_name + '-mask-' + mask_ame)
            mask_file.save(mask_path)
            masked_path = os.path.join(UPLOAD_SUBDIR, f_name + '-masked-' + mask_ame)
            add_mask2image_binary(path, mask_path, masked_path)

            path = masked_path
        # print(request.files['mask'])
        # mask = request.form['mask']
    except OSError as event:
        print(event)

    to_file_path = os.path.join(UPLOAD_SUBDIR, f_name + '-kou' + ex_name)

    if os.path.isfile(to_file_path):
        file_base64 = base64.b64encode(open(to_file_path, 'rb').read())
        result_text = {"statusCode": 200,
                       "filename": to_file_path, "message": "文件抠图成功,利用历史", "code": file_base64.decode()}
        response = make_response(jsonify(result_text))
        response.headers['Access-Control-Allow-Origin'] = '*'
        response.headers['Access-Control-Allow-Methods'] = 'OPTIONS,HEAD,GET,POST'
        response.headers['Access-Control-Allow-Headers'] = 'x-requested-with'
        return response

    if not os.path.isfile(path):
        result_text = {"statusCode": 500,
                       "filename": 'none', "message": "主文件不存在"}
        response = make_response(jsonify(result_text))
        return response

    try:
        os.system("conda run -n sam rembg i {} {}".format(path, to_file_path))
        file_base64 = base64.b64encode(open(to_file_path, 'rb').read())
    except OSError as event:
        print(event)

    result_text = {"statusCode": 200,
                   "filename": to_file_path, "message": "文件抠图成功", "code": file_base64.decode()}
    response = make_response(jsonify(result_text))
    response.headers['Access-Control-Allow-Origin'] = '*'
    response.headers['Access-Control-Allow-Methods'] = 'OPTIONS,HEAD,GET,POST'
    response.headers['Access-Control-Allow-Headers'] = 'x-requested-with'

    return response

def add_mask2image_binary(images_path, masks_path, masked_path):
    # Add binary masks to images
    img_path = images_path
    img = cv2.imread(img_path)
    # mask是.png格式的，image是.jpg格式的
    mask_path = masks_path
    # 将彩色mask以二值图像形式读取
    mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
    # 将image的相素值和mask像素值相加得到结果
    masked = cv2.add(img, np.zeros(np.shape(img), dtype=np.uint8), mask=mask)
    cv2.imwrite(masked_path, masked)

@app.route('/koufile', methods=['POST'])
def koufile():
    print(request)
    print(request.files['file'])
    # input标签中的name的属性值
    f = request.files['file']

    # 拼接地址，上传地址，f.filename：直接获取文件名
    f.save(os.path.join('./uploads', f.filename))
    path = os.path.join('./uploads', f.filename)
    # time.sleep(0.2)
    image = cut_image(path)  # Image.open(path)
    image.save(path)

    print("rembg i {} {}".format(path, path+'.to.' + f.filename))
    try:
        os.system("conda run -n sam rembg i {} {}".format(path,
                  path+'.to.' + f.filename))
    except OSError as e:
        print(e)

    result_text = {"statusCode": 200, "filename": f.filename +
                   '.to.' + f.filename, "message": "文件抠图成功"}
    response = make_response(jsonify(result_text))
    response.headers['Access-Control-Allow-Origin'] = '*'
    response.headers['Access-Control-Allow-Methods'] = 'OPTIONS,HEAD,GET,POST'
    response.headers['Access-Control-Allow-Headers'] = 'x-requested-with'

    return response


@app.route('/textocr', methods=['POST'])
def textocr():
    print(request)
    print(request.files['file'])
    # input标签中的name的属性值
    f = request.files['file']

    # 拼接地址，上传地址，f.filename：直接获取文件名
    f.save(os.path.join('./uploads', f.filename))
    path = os.path.join('./uploads', f.filename)

    # images = [path]
    # results = []
    # with PyTessBaseAPI() as api:
    #     for img in images:
    #         api.SetImageFile(img)
    #         re = {"texts": api.GetUTF8Text(), "poses": api.AllWordConfidences()}
    #         results.append(re)

    image = cut_image(path)  # Image.open(path)
    results = ''
    # with PyTessBaseAPI() as api:
    #     api.SetImage(image)
    #     boxes = api.GetComponentImages(RIL.TEXTLINE, True)
    #     print('Found {} textline image components.'.format(len(boxes)))
    #     for i, (im, box, _, _) in enumerate(boxes):
    #         # im is a PIL image object
    #         # box is a dict with x, y, w and h keys
    #         api.SetRectangle(box['x'], box['y'], box['w'], box['h'])
    #         ocrResult = api.GetUTF8Text()
    #         conf = api.MeanTextConf()
    #         results.append({"index": i, "conf": conf, "ocrResult": ocrResult, "box": box})
    #         print(u"Box[{0}]: x={x}, y={y}, w={w}, h={h}, "
    #               "confidence: {1}, text: {2}".format(i, conf, ocrResult, **box))
    results = tesserocr.image_to_text(image, lang='eng+chi_sim')

    result_text = {"statusCode": 200, "results": results, "message": "识别成功"}
    response = make_response(jsonify(result_text))
    response.headers['Access-Control-Allow-Origin'] = '*'
    response.headers['Access-Control-Allow-Methods'] = 'OPTIONS,HEAD,GET,POST'
    response.headers['Access-Control-Allow-Headers'] = 'x-requested-with'

    return response


@app.route('/uploader', methods=['POST', 'GET'])
def uploader():
    """
        文件上传
    """
    if request.method == 'POST':
        # input标签中的name的属性值
        f = request.files['file']
        fname, exname = os.path.splitext(f.filename)
        filename = str(uuid.uuid1()) + exname
        # 拼接地址，上传地址，f.filename：直接获取文件名
        f.save(os.path.join(UPLOAD_SUBDIR, filename))
        # 输出上传的文件名
        result_text = {"statusCode": 200,
                       "filename": filename, "message": "文件上传成功"}
        response = make_response(jsonify(result_text))
        response.headers['Access-Control-Allow-Origin'] = '*'
        response.headers['Access-Control-Allow-Methods'] = 'OPTIONS,HEAD,GET,POST'
        response.headers['Access-Control-Allow-Headers'] = 'x-requested-with'
        print(request.files, filename)
        return response
    else:
        return send_from_directory('static', 'home.html')


@app.route('/nsamcut', methods=['POST'])
def nsamcut():
    """
        上传图片, 调用sam切图
    """
    if request.method == 'POST':
        try:
            filename = request.form['filename']
            path = os.path.join(UPLOAD_SUBDIR, filename)
            f_name, ex_name = os.path.splitext(filename)
        except OSError as event:
            print(event)

        file_dir = os.path.join(UPLOAD_SUBDIR, 'samcut-dir')
        if not os.path.exists(file_dir):
            try:
                os.makedirs(file_dir)
            except OSError as e:
                print(e)

        ori_file_path = path
        parentDir = os.path.dirname(os.path.realpath(__file__))
        parent_parent = os.path.dirname(parentDir)
        sam_dir = os.path.join(parent_parent, 'segment-anything')

        command = ""
        command = command + "conda run -n sam python " + sam_dir + "/scripts/amg.py --checkpoint "+sam_dir + \
            "/models/sam_vit_h_4b8939.pth --model-type vit_h --input " + \
            ori_file_path+" --output "+file_dir

        print(command)
        comstrs = ""
        try:
            comstrs = os.system(command)
        except OSError as e:
            print(e)

        file_list = []
        fname, exname = os.path.splitext(filename)
        for file in os.listdir(os.path.join(file_dir, fname)):
            file_list.append(file)

        result_text = {"statusCode": 200, "command": command, "mask": file_list,
                       "filename": filename, "message": "文件上传成功", "result": comstrs}
        response = make_response(jsonify(result_text))
        response.headers['Access-Control-Allow-Origin'] = '*'
        response.headers['Access-Control-Allow-Methods'] = 'OPTIONS,HEAD,GET,POST'
        response.headers['Access-Control-Allow-Headers'] = 'x-requested-with'
        return response
    else:
        return send_from_directory('static', 'home.html')

@app.route('/samcut', methods=['POST'])
def samcut():
    """
        上传图片, 调用sam切图
    """
    if request.method == 'POST':
        file_dir = os.path.join(os.path.dirname(
            os.path.realpath(__file__)), 'uploads', 'samcut-dir')
        f = request.files['file']
        if not os.path.exists(file_dir):
            try:
                os.makedirs(file_dir)
            except OSError as e:
                print(e)
        f.save(os.path.join(file_dir, f.filename))

        ori_file_path = os.path.join(file_dir, f.filename)
        parentDir = os.path.dirname(os.path.realpath(__file__))
        parent_parent = os.path.dirname(parentDir)
        sam_dir = os.path.join(parent_parent, 'segment-anything')

        command = ""
        command = command + "conda run -n sam python " + sam_dir + "/scripts/amg.py --checkpoint "+sam_dir + \
            "/models/sam_vit_h_4b8939.pth --model-type vit_h --input " + \
            ori_file_path+" --output "+file_dir

        print(command)
        comstrs = ""
        try:
            comstrs = os.system(command)
        except OSError as e:
            print(e)

        file_list = []
        fname, exname = os.path.splitext(f.filename)
        for file in os.listdir(os.path.join(file_dir, fname)):
            file_list.append(file)

        result_text = {"statusCode": 200, "command": command, "mask": file_list,
                       "filename": f.filename, "message": "文件上传成功", "result": comstrs}
        response = make_response(jsonify(result_text))
        response.headers['Access-Control-Allow-Origin'] = '*'
        response.headers['Access-Control-Allow-Methods'] = 'OPTIONS,HEAD,GET,POST'
        response.headers['Access-Control-Allow-Headers'] = 'x-requested-with'
        return response
    else:
        return send_from_directory('static', 'home.html')


@app.route('/chachu', methods=['POST'])
def chachu():
    """
        上传图片以及 mask, 擦除mask指定元素
    """
    print('檫除...')
    if request.method == 'POST':
        dir = str(uuid.uuid1())
        # input标签中的name的属性值
        f = request.files['file']
        fmask = request.files['filemask']
        try:
            os.makedirs(os.path.join('./uploads', dir))
        except OSError as e:
            print(e)
        ori_filename = os.path.splitext(f.filename)[0] + '.png'
        # 拼接地址，上传地址，f.filename：直接获取文件名
        f.save(os.path.join('./uploads', dir, ori_filename))
        fmask.save(os.path.join('./uploads', dir, fmask.filename))

        curDir = os.path.realpath(__file__)
        parentDir = os.path.dirname(os.path.realpath(__file__))
        parent_parent = os.path.dirname(parentDir)
        # command = "conda activate lama && "
        command = ""
        command = command + "conda run -n lama export TORCH_HOME=" + \
            os.path.join(parent_parent, 'lama') + " && export PYTHONPATH=" + \
            os.path.join(parent_parent, "lama") + " && "
        command = command + "conda run -n lama python3 " + \
            os.path.join(parent_parent, "lama", "bin", "predict.py") + " model.path=" + \
            os.path.join(parent_parent, "lama", "LaMa_models", "big-lama")
        command = command + " indir=" + \
            os.path.join(parentDir, "uploads", dir)+" outdir=" + \
            os.path.join(parentDir, "uploads", dir+"-output")
        # command = command + " device=cpu"
        print(command)
        try:
            comstrs = os.system(command)
        except OSError as e:
            print(e)
        # 输出上传的文件名
        result_text = {"statusCode": 200, "dir": dir, "commands": comstrs,
                       "filename": ori_filename, "filemask": fmask.filename,  "message": "文件上传成功"}
        response = make_response(jsonify(result_text))
        response.headers['Access-Control-Allow-Origin'] = '*'
        response.headers['Access-Control-Allow-Methods'] = 'OPTIONS,HEAD,GET,POST'
        response.headers['Access-Control-Allow-Headers'] = 'x-requested-with'
        # print(request.files, f.filename)
        return response
    else:
        return send_from_directory('static', 'index.html')


@app.route(r'/assets/<path:path>', methods=['GET'])
def send_static(path):
    print(path)
    filepath = path is None and 'home' or path
    fpath = os.path.join('static', 'assets')
    print(fpath)
    return send_from_directory(fpath, filepath)


@app.route(r'/uploads/<path:path>', methods=['GET'])
def send_uploads(path):
    filepath = path
    return send_from_directory('uploads', filepath)


@app.route(r'/images/<path:path>', methods=['GET'])
def send_images(path):
    filepath = path
    fpath = os.path.join('uploads', 'samcut-dir')
    return send_from_directory(fpath, filepath)


def cut_image(path):
    # Read input image, and convert to NumPy array.
    # img is 1080 rows by 1920 cols and 4 color channels, the 4'th channel is alpha.
    img = np.array(Image.open(path))

    # Find indices of non-transparent pixels (indices where alpha channel value is above zero).
    idx = np.where(img[:, :, 3] > 0)

    # Get minimum and maximum index in both axes (top left corner and bottom right corner)
    x0, y0, x1, y1 = idx[1].min() if len(
        idx[1]) else 0, idx[0].min(), idx[1].max(), idx[0].max()

    # Crop rectangle and convert to Image
    out = Image.fromarray(img[y0:y1 + 1, x0:x1 + 1, :])

    # Save the result (RGBA color format).
    out.save('uploads/tmp.png')
    return out


if __name__ == "__main__":
    app.secret_key = 'super secret key'
    app.config['SESSION_TYPE'] = 'filesystem'
    app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

    CORS(app, resources=r'/*')

    sess.init_app(app)

    app.debug = True
    app.run(host='0.0.0.0', port=5001, debug=True)
